import os
from typing import Annotated, List, Dict, Any
from langchain_deepseek import ChatDeepSeek
from langchain_tavily import TavilySearch
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from typing_extensions import TypedDict
from dotenv import load_dotenv
import re
import json
from loguru import logger
import asyncio
from concurrent.futures import ThreadPoolExecutor
import time

# 加载环境变量
load_dotenv("./static/config.env")
DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
TAVILY_API_KEY = os.getenv("TAVILY_KEY")

# 初始化工具和模型
tavily_search_tool = TavilySearch(max_results=2, tavily_api_key=TAVILY_API_KEY)  # 减少搜索结果数量

llm = ChatDeepSeek(model="deepseek-chat", temperature=0.7, api_key=DEEPSEEK_API_KEY, request_timeout=30)

class QAState(TypedDict):
    messages: Annotated[list, add_messages]
    questions: List[str]
    results: Dict[str, Any]


# 使用线程池来并行执行搜索
executor = ThreadPoolExecutor(max_workers=5)


async def async_search(query: str):
    """异步搜索函数"""
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(executor, tavily_search_tool.invoke, {"query": query})


async def search_node(state: QAState):
    """异步搜索节点：并行搜索所有问题"""
    questions = state["questions"]
    search_results = {}

    logger.info(f"开始并行搜索 {len(questions)} 个问题")
    start_time = time.time()

    # 并行执行所有搜索
    search_tasks = [async_search(question) for question in questions]
    search_responses = await asyncio.gather(*search_tasks, return_exceptions=True)

    # 处理搜索结果
    for question, response in zip(questions, search_responses):
        if isinstance(response, Exception):
            search_results[question] = f"搜索出错: {str(response)}"
        else:
            search_results[question] = response

    elapsed_time = time.time() - start_time
    logger.info(f"并行搜索完成，耗时: {elapsed_time:.2f}秒")

    return {"messages": [AIMessage(content=f"已完成搜索，共处理{len(questions)}个问题")], "results": search_results}


async def answer_node(state: QAState):
    """异步回答节点：基于搜索结果生成结构化答案"""
    questions = state["questions"]
    search_results = state.get("results", {})
    structured_answers = {}

    logger.info(f"开始并行生成 {len(questions)} 个问题的结构化答案")
    start_time = time.time()

    # 构建所有提示词
    prompts = []
    for question in questions:
        search_info = search_results.get(question, [])

        # 构建结构化回答提示词
        system_prompt = """
            你是一个专业的知识解答助手，能够基于搜索结果为用户问题提供准确、全面的回答，并以结构化格式输出。
            请遵循以下要求：
            1. 基于提供的搜索结果回答问题
            2. 回答要有条理，使用项目符号或编号列表
            3. 如果搜索结果不足以回答问题，请说明
            4. 以严格的JSON格式返回结果，包含以下字段：
               - title: 每一个回答的主要概括性总结（简洁概括回答核心），必须要有值，你可以自由发挥总结网页，最好保持在15字以内
               - content: 回答的正文内容
               - citations: 搜索结果中返回的引用链接列表，包含url和title字段
            
            示例格式：
            {
              "title": "示例标题",
              "content": "示例内容",
              "citations": [
                {
                  "url": "https://example.com",
                  "title": "示例链接标题"
                }
              ]
            }
            重要规则: 
                - 只返回纯JSON字符串，不要使用任何Markdown代码块包装
                - 只返回JSON字符串，不要添加任何额外说明
            """

        user_prompt = f"""
            问题：{question}
            搜索结果：
            {search_info}
            请基于上述搜索结果回答问题，并严格按照指定的JSON字符串格式返回结果
            """

        prompts.append((question, system_prompt, user_prompt))

    # 并行调用LLM生成结构化答案
    async def generate_structured_answer(question, system_prompt, user_prompt):
        try:
            response = await llm.ainvoke([
                SystemMessage(content=system_prompt),
                HumanMessage(content=user_prompt)
            ])

            # 尝试解析JSON
            try:
                structured_answer = json.loads(response.content)
            except json.JSONDecodeError:
                # 如果LLM没有返回有效的JSON，构造默认结构
                structured_answer = {
                    "title": question,
                    "content": response.content if hasattr(response, 'content') else str(response),
                    "citations": _extract_citations_from_content(response.content if hasattr(response, 'content') else str(response))
                }

            return question, structured_answer
        except Exception as e:
            # 出错时返回默认结构
            default_answer = {
                "title": question,
                "content": f"生成答案时出错: {str(e)}",
                "citations": []
            }
            return question, default_answer

    # 创建并执行所有答案生成任务
    answer_tasks = [generate_structured_answer(q, sys_prompt, user_prompt) for q, sys_prompt, user_prompt in prompts]
    results = await asyncio.gather(*answer_tasks)

    # 收集结果
    for question, structured_answer in results:
        # 确保返回的数据结构完整
        if not isinstance(structured_answer, dict):
            structured_answers[question] = {
                "title": question,
                "content": str(structured_answer),
                "citations": []
            }
        else:
            structured_answers[question] = {
                "title": structured_answer.get("title", question),
                "content": structured_answer.get("content", ""),
                "citations": structured_answer.get("citations", [])
            }

    elapsed_time = time.time() - start_time
    logger.info(f"并行生成结构化答案完成，耗时: {elapsed_time:.2f}秒")

    return {"messages": [AIMessage(content="已完成结构化回答生成")], "results": structured_answers}


def _extract_citations_from_content(content: str) -> List[Dict[str, str]]:
    """从内容中提取引用链接"""
    # 提取可能的链接
    urls = re.findall(r'https?://[^\s<>"]+|www\.[^\s<>"]+', content)
    citations = [{"url": url, "title": f"引用链接 {i + 1}"} for i, url in enumerate(urls)]
    return citations


# 构建图
graph_builder = StateGraph(QAState)

# 使用异步搜索节点和异步回答节点
graph_builder.add_node("search", search_node)
graph_builder.add_node("answer", answer_node)

graph_builder.add_edge(START, "search")
graph_builder.add_edge("search", "answer")
graph_builder.add_edge("answer", END)

# 编译图
graph = graph_builder.compile()


async def answer_questions_with_sources(questions: List[str]) -> Dict[str, Any]:
    """主函数：接收问题列表，返回带引用的答案"""
    initial_state = {
        "messages": [],
        "questions": questions,
        "results": {}
    }

    # 执行图（使用异步执行来支持异步搜索节点）
    final_state = await graph.ainvoke(initial_state)

    logger.info("返回的答案结果: {}".format(str(final_state["results"])[:1000]))  # 限制长度避免日志过长
    return final_state["results"]
